# MCP Server工具 - MCP Server 包含多类工具
import os
from dotenv import load_dotenv
import json
import pymysql
from mcp.server.fastmcp import FastMCP


# 加载环境变量
load_dotenv(override=True)
host = os.getenv('HOST')
user = os.getenv('USER')
mysql_pw = os.getenv('MYSQL_PW')
db = os.getenv('DB_NAME')
port = os.getenv('PORT')
# 创建MCP
mcp = FastMCP("AgentMySQLServer", port=3366)


# Tools
@mcp.tool()
async def sql_inquire(sql_query: str) -> str:
    """
    当用户需要进行数据库查询工作时，请调用该函数。
    该函数用于在指定MySQL服务器上运行一段SQL代码，完成数据查询相关工作，
    并且当前函数是使用pymsql连接MySQL数据库。
    本函数只负责运行SQL代码并进行数据查询。
    :param sql_query: 字符串形式的SQL查询语句，用于执行对MySQL数据库中insight_general_table表进行查询，并获得相关信息
    :return：sql_query在MySQL中的运行结果。
    """
    # 创建连接
    connection = pymysql.connect(
        host=host,
        user=user,
        passwd=mysql_pw,
        db=db,
        port=int(port),
        charset='utf8'
    )
    try:
        with connection.cursor() as cursor:
            cursor.execute(sql_query)
            results = cursor.fetchall()
            # print("SQL 查询已成功执行，正在整理结果...")
    finally:
        connection.close()

    # 将结果以 JSON 字符串形式返回
    return json.dumps(results, ensure_ascii=False)


if __name__ == "__main__":
    mcp.run(transport="streamable-http")